# tts/tts_client.py
# TTS HTTP客户端，负责与TTS服务器的通信

import requests
import time
from typing import Iterator, Optional


class TTSClient:
    """TTS HTTP客户端，处理与TTS服务器的所有通信"""
    
    def __init__(self, api_host: str, tts_url: str, control_url: str):
        self.api_host = api_host
        self.tts_url = tts_url
        self.control_url = control_url
        self.is_available = True  # TTS服务可用性状态
    
    def check_connection(self) -> bool:
        """检测TTS服务连接状态"""
        try:
            # 尝试连接控制接口
            response = requests.get(f"{self.control_url}?command=status", timeout=5)
            if response.status_code == 200:
                print("✅ TTS服务连接正常")
                self.is_available = True
                return True
            else:
                print(f"❌ TTS服务连接异常 (状态码: {response.status_code})")
                self.is_available = False
                return False
        except requests.exceptions.RequestException as e:
            print(f"❌ TTS服务连接失败: {e}")
            self.is_available = False
            return False
    
    def switch_models(self, gpt_path: str, sovits_path: str) -> bool:
        """通过 API 接口切换 GPT 和 SoVITS 模型权重"""
        print("--- 1. 正在切换 GPT 模型...")
        encoded_gpt_path = requests.utils.quote(gpt_path)
        gpt_url = f"{self.api_host}/set_gpt_weights?weights_path={encoded_gpt_path}"
        
        try:
            response_gpt = requests.get(gpt_url, timeout=60)
            if response_gpt.status_code == 200:
                print(f"✅ GPT 模型切换成功: {gpt_path}")
            else:
                print(f"❌ GPT 模型切换失败 ({response_gpt.status_code}): {response_gpt.text}")
                return False
        except requests.exceptions.RequestException as e:
            print(f"❌ 切换 GPT 模型时连接失败: {e}")
            self.is_available = False
            return False

        print("--- 2. 正在切换 SoVITS 模型...")
        encoded_sovits_path = requests.utils.quote(sovits_path)
        sovits_url = f"{self.api_host}/set_sovits_weights?weights_path={encoded_sovits_path}"
        
        try:
            response_sovits = requests.get(sovits_url, timeout=60)
            if response_sovits.status_code == 200:
                print(f"✅ SoVITS 模型切换成功: {sovits_path}")
                return True
            else:
                print(f"❌ SoVITS 模型切换失败 ({response_sovits.status_code}): {response_sovits.text}")
                return False
        except requests.exceptions.RequestException as e:
            print(f"❌ 切换 SoVITS 模型时连接失败: {e}")
            self.is_available = False
            return False

    def preload_tts(self, params: dict, preload_text: str) -> bool:
        """通过一个简短的合成请求来预热模型"""
        print(f"\n--- 正在进行模型预热（Preload）：'{preload_text}' ---")
        preload_params = params.copy()
        preload_params["text"] = preload_text 
        
        start_time = time.time()
        try:
            response = requests.get(self.tts_url, params=preload_params, timeout=120)
            response.raise_for_status()
            
            if len(response.content) > 44:
                duration = time.time() - start_time
                print(f"✅ 模型预热成功。耗时: {duration:.3f} 秒。")
                return True
            else:
                print(f"❌ 模型预热失败：API返回的数据过短。")
                return False
                
        except requests.exceptions.RequestException as e:
            print(f"❌ 模型预热时网络请求异常: {e}")
            self.is_available = False
            return False
        except Exception as e:
            print(f"❌ 模型预热时发生未知错误: {e}")
            self.is_available = False
            return False

    def get_tts_audio_stream(self, text: str, tts_params: dict, chunk_size: int, 
                           stop_event=None) -> Iterator[bytes]:
        """获取TTS音频流，简化处理逻辑"""
        # 如果TTS服务不可用，直接返回空迭代器
        if not self.is_available:
            return
        
        params = tts_params.copy()
        params["text"] = text
        
        response = None 
        try:
            response = requests.get(
                self.tts_url, 
                params=params, 
                stream=True, 
                timeout=30
            )
            
            is_hard_interrupted = response.status_code >= 400 and stop_event and stop_event.is_set()
            if is_hard_interrupted:
                print(f"\n[TTS] TTS服务器因硬中断而提前终止流（状态码: {response.status_code}）。")
                return 
            
            response.raise_for_status()

            first_chunk = True
            
            for chunk in response.iter_content(chunk_size):
                # 更频繁地检查停止事件
                if stop_event and stop_event.is_set():
                    print("\n[TTS] TTS合成被软中断，停止接收流...")
                    return
                if chunk:
                    if first_chunk:
                        # 简单处理：如果首块只有WAV头部，等待下一个块
                        if len(chunk) <= 44:
                            continue
                        else:
                            # 首块包含音频数据，跳过头部
                            if chunk[:4] == b'RIFF':
                                yield chunk[44:]
                            else:
                                yield chunk
                        first_chunk = False
                    else:
                        yield chunk
                        
        except requests.exceptions.RequestException as e:
            if not stop_event or not stop_event.is_set():
                print(f"\n❌ TTS 网络请求异常: {e}")
        except Exception as e:
            print(f"\n❌ 处理TTS流时发生未知错误: {e}")
        finally:
            if response:
                response.close()
    
    def send_stop_command(self):
        """发送停止命令到TTS服务器"""
        try:
            response = requests.get(f"{self.control_url}?command=stop", timeout=2)
            if response.status_code == 200:
                print("\n[系统] TTS服务器硬中断命令发送成功。")
            else:
                print(f"\n[系统] TTS硬中断命令发送，API返回状态码: {response.status_code}")
        except requests.exceptions.RequestException as e:
            print(f"\n❌ [系统] 发送 TTS 硬中断命令时网络请求异常: {e}")
